import pymysql
import csv
import os

def backup_tables_to_csv(folder_path, tables):
    connection = pymysql.connect(
        host='localhost',
        user='root',
        password='123456',
        database='demo01'
    )
    try:
        cursor = connection.cursor()
        for table in tables:
            # 检查备份文件夹是否存在，不存在则创建
            if not os.path.exists(folder_path):
                os.makedirs(folder_path)
            backup_file_name = f"{table}_backup.csv"
            backup_file_path = os.path.join(folder_path, backup_file_name)
            with open(backup_file_path, 'w', newline='') as csv_file:
                writer = csv.writer(csv_file)
                # 获取表的列名
                cursor.execute(f"SHOW COLUMNS FROM {table}")
                columns = [column[0] for column in cursor.fetchall()]
                writer.writerow(columns)
                # 获取表的数据
                cursor.execute(f"SELECT * FROM {table}")
                rows = cursor.fetchall()
                for row in rows:
                    writer.writerow(row)
    except pymysql.Error as e:
        print(f"Error during backup: {e}")
    finally:
        connection.close()

def restore_tables_from_csv(folder_path, tables):
    connection = pymysql.connect(
        host='localhost',
        user='root',
        password='123456',
        database='demo01'
    )
    try:
        cursor = connection.cursor()
        for table in tables:
            backup_file_name = f"{table}_backup.csv"
            backup_file_path = os.path.join(folder_path, backup_file_name)
            if os.path.exists(backup_file_path):
                with open(backup_file_path, 'r', newline='') as csv_file:
                    reader = csv.reader(csv_file)
                    columns = next(reader)
                    for row in reader:
                        values = []
                        for value in row:
                            if isinstance(value, str):
                                value = f'"{value}"'
                            values.append(value)
                        insert_statement = f"INSERT INTO {table} ({', '.join(columns)}) VALUES ({', '.join(values)})"
                        try:
                            cursor.execute(insert_statement)
                        except pymysql.Error as e:
                            print(f"Error inserting data into {table}: {e}")
                connection.commit()
    except pymysql.Error as e:
        print(f"Error during restoration: {e}")
    finally:
        connection.close()

if __name__ == "__main__":
    backup_folder = "backup"
    tables_to_backup = ["users", "documents"]
    # backup_tables_to_csv(backup_folder, tables_to_backup)
    restore_tables_from_csv(backup_folder, tables_to_backup)